from dotenv import load_dotenv, find_dotenv
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage

_ = load_dotenv(find_dotenv())

loader = PyPDFLoader("llama2.pdf")
pages = loader.load_and_split()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=200,
    chunk_overlap=100,
    length_function=len,
    add_start_index=True,
)
paragraphs = text_splitter.create_documents([page.page_content for page in pages[:4]])

embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
db = Chroma.from_documents(paragraphs, embeddings)

# 检索 top-3 结果
retriever = db.as_retriever(search_kwargs={"k": 3})

query = "llama2有多少参数"

docs = retriever.invoke(query)

messages = []

for doc in docs:
    print(doc)
    print("-------------------")
    messages.append(AIMessage(content=doc.page_content))
    
messages.append(HumanMessage(content=query))

llm = ChatOpenAI(model="gpt-4o")
response = llm.invoke(messages)
print(response.content)
    
